//
//  equations.cpp
//  EMT_V2_Crisis
//
//  Created by Duong Ngo on 6/5/17.
//  Copyright © 2017 Duong Ngo. All rights reserved.
//

#include "equations.hpp"
#include "global.hpp"
#include "mapping.hpp"

#include "math.h"

void set_up_equations_at_time_t (const int& t, const std::vector<double>& ss, const Number* x, Number* g)
{
    // Mapping variable from x to s
    variables s;
    find_variable_s_at_time_t(t, x, s);
    
    // Position of equation
    int pg = (t-1)*var_foc;
    
    // List of equations
    //Bankers
    g[pg+0]= s.gamma*s.cb- 1;
    g[pg+1]= s.gamma- (beta_b*s.Rf*s.gamma_f*s.ip_f+ pe*pow(fmax(s.muc,0),2));
    g[pg+2]= s.gamma- (beta_b*s.Rm*s.gamma_f*s.ip_f+ pe*pow(fmax(s.muc,0),2)+ varphi*s.mur);
    g[pg+3]= s.gamma- (beta_b*Rn*s.gamma_f*s.ip_f+ pe*pow(fmax(s.muc,0),2)+ s.mur);
    g[pg+4]= (s.ql+thetaa)*s.gamma- (beta_b*(deltab+ deltab*s.ql_f)*s.gamma_f*s.ip_f+ (1-vec_kappa[t])*pe*pow(fmax(s.muc,0),2));
    g[pg+5]= s.n_p*s.ip+ s.tau- s.n;
    g[pg+6]= s.m- (s.Rm_p*s.m_p*s.ip+ s.ql*s.s+ thetaa*s.bh - deltab*s.bh_p*s.ip+ s.cb+ s.tau- (Rn-1)*s.n_p*s.ip);
    g[pg+7]= s.n- varphi*s.m;
    g[pg+8]= s.n+ (1-vec_kappa[t])*s.bh- s.m+ s.muc;
    g[pg+9]= s.bh- deltab*s.bh_p*s.ip- s.s;
    
    //Households
    g[pg+10]= 1/s.ch- (pow(fmax(s.etaz,0),2) + s.lama);
    g[pg+11]= s.ch*s.lamb- 1;
    g[pg+12]= s.lama- beta_h*s.Rm*s.lamb_f*s.ip_f;
    g[pg+13]= s.ql*s.lamb- ( beta_h*(deltab+ deltab*s.ql_f)*s.lamb_f*s.ip_f+ pe*pow(fmax(s.etab,0),2) );
    g[pg+14]= s.lamb- ( beta_h*(1-deltak)*s.lamb_f+ beta_h*alphaa*s.pm_f*s.lama_f*s.y_f/s.k );
    g[pg+15]= s.Rm_p*s.m_p*s.ip+ s.ql*s.s- s.ch- s.i- deltab*s.bh_p*s.ip- pow(fmax(-s.etaz,0),2);
    g[pg+16]= bh_limit- s.bh+ s.etab;
    
    //Inflation and Output
    g[pg+17]= 1- iota*(s.pie-1)*s.pie+ iota*beta_h*s.lama_f/s.lama*(s.pie_f-1)*s.pie_f*s.y_f/s.y- (1-s.pm)*epsilon;
    g[pg+18]= s.y- pow(s.k_p, alphaa)*pow(s.l, 1-alphaa);
    
    //Central bank
    g[pg+19]= s.u- (1/beta_b*pow(s.pie_f, phipie)* pow(s.y/ss[original_map.at("s.y")], phiy)*exp(shock_rf[t]));
    
    //Market clearing
    g[pg+20]= s.y- (s.cb+ s.ch+ s.i+ thetaa*s.bh+ iota/2*(s.pie-1)*(s.pie-1)*s.y);
    g[pg+21]= s.k- ((1-deltak)*s.k_p+ s.i);
    g[pg+22]= s.pie*s.ip- 1;

    // Soft max constraint
    if (s.u > Rf_min){
        double texp= exp(s_max*(Rf_min- s.u));
        g[pg+23]= s.Rf- (s.u+ log(1+texp)/s_max);
    }
    else {
        double texp= exp(s_max*(s.u- Rf_min));
        g[pg+23]= s.Rf- (Rf_min+ log(1+texp)/s_max);
    }

    
    // Labor
    g[pg+24]= chi*pow(s.l, nu+1)- (1-alphaa)*s.pm*s.y*s.lama;
    
    
}

//------------------------------------------------------------------------------------------------------
//------------------------------------------------------------------------------------------------------
//------------------------------------------------------------------------------------------------------

void set_up_initial_conditon_at_time_zero (const std::vector<double>& ss, const Number* x, Number* g)
{
    int pg= (T-1)*var_foc;
    for (int i=0; i<var_foc; ++i){
        if (var_role[i] < dim_state){
            g[pg+i] = x[i]- starting_state[var_role[i]];
        }
        else {
            g[pg+i] = x[i]- ss[i];
        }
    }
}

void set_up_terminal_condition_at_final_time (const std::vector<double>& ss, const Number* x, Number* g)
{
    int pg= T*var_foc;
    int px= T*var_foc;
    for (int i=0; i<var_foc; ++i){
        g[pg+i] = x[px+i]- ss[i];
    }
}

//------------------------------------------------------------------------------------------------------
//------------------------------------------------------------------------------------------------------
//------------------------------------------------------------------------------------------------------

void summarize_all_equations (const std::vector<double>& ss, const Number* x, Number* g)
{
    for (int t=1; t<T; ++t){
        set_up_equations_at_time_t(t, ss, x, g);
    }
    set_up_initial_conditon_at_time_zero(ss, x, g);
    set_up_terminal_condition_at_final_time(ss, x, g);
}
















